
"""
智能问数查询改写
"""
from core.state import AnalyzeState
from langgraph.types import Command
from langgraph.config import get_stream_writer
from core.llm import rewrite_query_analyzer
import concurrent.futures
from loguru import logger


async def rewrite_query_node(state: AnalyzeState):

    writer = get_stream_writer()
    writer("智能问数查询改写")
    logger.info("智能问数查询改写")

    schemas = state.get('schemas')
    query = state.get('query')
    rewrites = state.get('rewrites')
    userinfo = state.get('userinfo')

    # 智能问数查询改写。
    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
        tasks = []
        for schema in schemas:
            tasks.append(executor.submit(rewrite_query_analyzer.invoke,
                                         userinfo.description, schema, query))

        for future in concurrent.futures.as_completed(tasks):
            if future.exception() is not None:
                continue
            result = future.result()
            if not isinstance(result, dict):
                continue
            relevant = result.get('relevant', False)
            if str(relevant).lower() != 'true':
                continue
            if len(result.get('rewrite', '')) > 2:
                rewrites.append(result['rewrite'])

    writer(rewrites)
    logger.info(f'改写的问题描述列表：{rewrites}')

    return Command(update={"rewrites": rewrites}, goto='biw_query')
